#!/usr/bin/env python3
"""
简易分布式缓存系统 - 单个节点实现
使用 FastAPI 提供 HTTP 接口，支持分布式 Key-Value 存储
"""

import json
import hashlib
from typing import Dict, Any, Optional
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
import uvicorn
import httpx
import os
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class CacheNode:
    """缓存节点类，负责本地存储和分布式协调"""
    
    def __init__(self, node_id: int, total_nodes: int = 3, base_port: int = 8000):
        self.node_id = node_id
        self.total_nodes = total_nodes
        self.base_port = base_port
        self.local_cache: Dict[str, Any] = {}
        self.port = base_port + node_id
        
    def _get_target_node(self, key: str) -> int:
        """使用哈希算法确定key应该存储在哪个节点"""
        hash_value = int(hashlib.md5(key.encode()).hexdigest(), 16)
        return hash_value % self.total_nodes
    
    def _get_node_url(self, node_id: int) -> str:
        """获取指定节点的HTTP地址"""
        port = self.base_port + node_id
        return f"http://localhost:{port}"
    
    async def _forward_request(self, method: str, node_id: int, key: str = None, value: Any = None) -> Any:
        """将请求转发到目标节点"""
        target_url = self._get_node_url(node_id)
        
        async with httpx.AsyncClient() as client:
            try:
                if method == "GET":
                    response = await client.get(f"{target_url}/{key}")
                    if response.status_code == 404:
                        raise HTTPException(status_code=404, detail="Key not found")
                    return response.json()
                elif method == "POST":
                    response = await client.post(target_url, json={key: value})
                    return response.json()
                elif method == "DELETE":
                    response = await client.delete(f"{target_url}/{key}")
                    return response.json()
            except httpx.RequestError as e:
                logger.error(f"Error forwarding request to node {node_id}: {e}")
                raise HTTPException(status_code=503, detail="Service unavailable")

# 创建FastAPI应用
app = FastAPI(title="分布式缓存节点", version="1.0.0")

# 获取节点配置
node_id = int(os.getenv("NODE_ID", "0"))
total_nodes = int(os.getenv("TOTAL_NODES", "3"))
base_port = int(os.getenv("BASE_PORT", "8000"))

# 初始化缓存节点
cache_node = CacheNode(node_id, total_nodes, base_port)

@app.post("/")
async def set_value(request: Request):
    """写入/更新缓存值"""
    try:
        data = await request.json()
        if not data or len(data) != 1:
            raise HTTPException(status_code=400, detail="Request body must contain exactly one key-value pair")
        
        key, value = next(iter(data.items()))
        target_node = cache_node._get_target_node(key)
        
        logger.info(f"Setting key '{key}' (target node: {target_node})")
        
        if target_node == node_id:
            # 本地存储
            cache_node.local_cache[key] = value
            return JSONResponse(content={key: value}, status_code=200)
        else:
            # 转发到目标节点
            result = await cache_node._forward_request("POST", target_node, key, value)
            return JSONResponse(content=result, status_code=200)
            
    except Exception as e:
        logger.error(f"Error setting value: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/{key}")
async def get_value(key: str):
    """读取缓存值"""
    target_node = cache_node._get_target_node(key)
    
    logger.info(f"Getting key '{key}' (target node: {target_node})")
    
    if target_node == node_id:
        # 本地读取
        if key in cache_node.local_cache:
            return JSONResponse(content={key: cache_node.local_cache[key]}, status_code=200)
        else:
            raise HTTPException(status_code=404, detail="Key not found")
    else:
        # 转发到目标节点
        result = await cache_node._forward_request("GET", target_node, key)
        return JSONResponse(content=result, status_code=200)

@app.delete("/{key}")
async def delete_value(key: str):
    """删除缓存值"""
    target_node = cache_node._get_target_node(key)
    
    logger.info(f"Deleting key '{key}' (target node: {target_node})")
    
    if target_node == node_id:
        # 本地删除
        if key in cache_node.local_cache:
            del cache_node.local_cache[key]
            return JSONResponse(content={"deleted": 1}, status_code=200)
        else:
            return JSONResponse(content={"deleted": 0}, status_code=200)
    else:
        # 转发到目标节点
        result = await cache_node._forward_request("DELETE", target_node, key)
        return JSONResponse(content=result, status_code=200)

@app.get("/")
async def health_check():
    """健康检查接口"""
    return JSONResponse(content={
        "status": "healthy",
        "node_id": node_id,
        "total_nodes": total_nodes,
        "cache_size": len(cache_node.local_cache)
    }, status_code=200)

if __name__ == "__main__":
    port = base_port + node_id
    logger.info(f"Starting cache node {node_id} on port {port}")
    uvicorn.run(app, host="0.0.0.0", port=port)
